#ifndef BASE_ROLLINGACCUMULATOR_H_
#define BASE_ROLLINGACCUMULATOR_H_

#include <vector>

#include "common.h"

namespace base {

	// RollingAccumulator stores and reports statistics
	// over N most recent samples.
	//
	// T is assumed to be an int, long, double or float.
	template<typename T>
	class RollingAccumulator {
	public:
		explicit RollingAccumulator(size_t max_count)
			: count_(0),
			next_index_(0),
			sum_(0.0),
			sum_2_(0.0),
			samples_(max_count) {
		}
		~RollingAccumulator() {
		}

		size_t max_count() const {
			return samples_.size();
		}

		size_t count() const {
			return count_;
		}

		void AddSample(T sample) {
			if (count_ == max_count()) {
				// Remove oldest sample.
				T sample_to_remove = samples_[next_index_];
				sum_ -= sample_to_remove;
				sum_2_ -= sample_to_remove * sample_to_remove;
			} else {
				// Increase count of samples.
				++count_;
			}
			// Add new sample.
			samples_[next_index_] = sample;
			sum_ += sample;
			sum_2_ += sample * sample;
			// Update next_index_.
			next_index_ = (next_index_ + 1) % max_count();
		}

		T ComputeSum() const {
			return static_cast<T>(sum_);
		}

		T ComputeMean() const {
			if (count_ == 0) {
				return static_cast<T>(0);
			}
			return static_cast<T>(sum_ / count_);
		}

		// O(n) time complexity.
		// Weights nth sample with weight (learning_rate)^n. Learning_rate should be
		// between (0.0, 1.0], otherwise the non-weighted mean is returned.
		T ComputeWeightedMean(double learning_rate) const {
			if (count_ < 1 || learning_rate <= 0.0 || learning_rate >= 1.0) {
				return ComputeMean();
			}
			double weighted_mean = 0.0;
			double current_weight = 1.0;
			double weight_sum = 0.0;
			const size_t max_size = max_count();
			for (size_t i = 0; i < count_; ++i) {
				current_weight *= learning_rate;
				weight_sum += current_weight;
				// Add max_size to prevent underflow.
				size_t index = (next_index_ + max_size - i - 1) % max_size;
				weighted_mean += current_weight * samples_[index];
			}
			return static_cast<T>(weighted_mean / weight_sum);
		}

		// Compute estimated variance.  Estimation is more accurate
		// as the number of samples grows.
		T ComputeVariance() const {
			if (count_ == 0) {
				return static_cast<T>(0);
			}
			// Var = E[x^2] - (E[x])^2
			double count_inv = 1.0 / count_;
			double mean_2 = sum_2_ * count_inv;
			double mean = sum_ * count_inv;
			return static_cast<T>(mean_2 - (mean * mean));
		}

	private:
		size_t count_;
		size_t next_index_;
		double sum_;    // Sum(x)
		double sum_2_;  // Sum(x*x)
		std::vector<T> samples_;

		DISALLOW_COPY_AND_ASSIGN(RollingAccumulator);
	};

}  // namespace base

#endif  // BASE_ROLLINGACCUMULATOR_H_
